package business;

import htsjdk.variant.variantcontext.Genotype;
import htsjdk.variant.variantcontext.VariantContext;

import java.io.BufferedWriter;
import java.io.File;
import java.io.FileWriter;
import java.io.IOException;
import java.util.HashMap;
import java.util.HashSet;
import java.util.TreeSet;

import model.ModelInputObj;
import model.SampleExpression;

public class InputObjGen {
	private String gtfPath;
	private String mmseqPath;
	private String doePath;
	private String vcfPath;
	private String vcfIndexPath;
	private String covPath;
	private int BUFFER = 100; //additional region to query above or below what the .gtf specifies
	
	public InputObjGen (String gtf, String mmseq, String doe, String vcf, String vcfIndex, String cov)
	{
		gtfPath = gtf;
		mmseqPath = mmseq;
		doePath = doe;
		vcfPath = vcf;
		vcfIndexPath = vcfIndex;
		covPath = cov;
	}
	
	public HashSet<ModelInputObj> gen ()
	{
		//load data from all files except mmseq
		HashSet<ModelInputObj> results = new HashSet<ModelInputObj>();
		MMSeqReader mmseq = new MMSeqReader(mmseqPath);
		String[] mmseqLine = mmseq.readLine();
		GTFReader gtf = new GTFReader(gtfPath);
		HashMap<String, Transcript> gtfData = gtf.parseTranscripts();
		VCFReader vcf = new VCFReader(vcfPath,vcfIndexPath);
		DOEReader doe = new DOEReader(doePath);
		HashMap<String, HashSet<String[]>> doeData = doe.readDOE();
		COVReader cov = new COVReader(covPath);
		HashMap<String, String[]> covData = cov.readCOV();
		while(mmseqLine != null) //loop through mmseqLines while loading all other data
		{
			Transcript transcript = gtfData.get(mmseqLine[0].split("\"")[1]);
			
			HashSet<VariantContext> variants = vcf.query(transcript.getChromosome(), transcript.getStart()-BUFFER,transcript.getStop()+BUFFER,transcript.getTranscriptID()); //returns VariantContext Objects within the region specified
			for (VariantContext v: variants)
			{
				Variant variant = new Variant(); //sets Variant data from the VC Object
				variant.setChrom(v.getChr());
				variant.setPos(v.getEnd());
				variant.setVariant_id(v.getID());
				variant.setRef(v.getReference().toString().substring(0,1)); 
				variant.setAlt(v.getAlternateAllele(0).toString());
				for(Genotype genotype: v.getGenotypes()) //loops through each individual data (represented by a Genotype obj) within the VC Object
				{
					Individual individual = new Individual(); //sets Individual data from the data within the genotype Object
					individual.setVariant(variant);
					individual.setSubject_id(genotype.getSampleName());
					individual.setGene_name(transcript.getGeneName());
					String genotypeCode = "";
					String allele1 = genotype.getAllele(0).toString(); //extracting the allele type
					String allele2 = genotype.getAllele(1).toString();
					if(allele1.length() == 1) //determining if the alleles are reference or alternate (ref alleles are marked with *, so they have 2 characters)
						genotypeCode += 1;
					else genotypeCode += 0;
					if(allele2.length() == 1)
						genotypeCode += 1;
					else genotypeCode += 0;
					individual.setGenotypeCode(genotypeCode);
					individual.setGenotype(allele1.substring(0,1) + allele2.substring(0,1));
					individual.setCovariates(covData.get(individual.getSubject_id())); //searches HashMap from covData with the subject_id
					
					HashSet<String[]> samples = doeData.get(genotype.getSampleName());
					HashMap<Enum<Tissue>, HashSet<SampleExpression>> tissueExp = 
							new HashMap<Enum<Tissue>, HashSet<SampleExpression>>(); //data structure containing tissue-sample pairs
					if(samples == null) //if the doe does not include the sampleName, skip this genotype object entirely
					{
						continue;
					}
					for(String[] sampleData: samples) //loop through all samples from this individual
					{
						Sample sample = new Sample(); //set sample data from the doe file
						sample.setSample_id(sampleData[5]);
						sample.setIndividual(individual);
						String tissue = sampleData[19];
						sample.setTissue(Tissue.valueOf(tissue.toLowerCase()));
						
						SampleExpression sampExp = new SampleExpression(); //set SampleExp obj from all previous data
						sampExp.setSample(sample);
						sampExp.setIndividual(individual);
						int sampleLoc = mmseq.getSampleNames().indexOf(sample.getSample_id());
						if(sampleLoc == -1) //if mmseq does not contain this sample, skip it
						{
							continue;
						}
						sampleLoc = sampleLoc * 8 + 1;
						sampExp.setMean_hap1((mmseqLine[sampleLoc])); //search mmseq for sample name and set the data
						sampExp.setVar_hap1((mmseqLine[sampleLoc+1]));
						sampExp.setMean_hap2((mmseqLine[sampleLoc+2]));
						sampExp.setVar_hap2((mmseqLine[sampleLoc+3]));
						sampExp.setMdiff((mmseqLine[sampleLoc+4]));
						sampExp.setVdiff((mmseqLine[sampleLoc+5]));
						sampExp.setUnique_hit_hap1((mmseqLine[sampleLoc+6]));
						sampExp.setUnique_hit_hap2((mmseqLine[sampleLoc+7]));
						if(tissueExp.get(sample.getTissue()) == null) 
							//if there is no previous HashSet for this tissue make a new one
						{
							HashSet<SampleExpression> hSet = new HashSet<SampleExpression>();
							hSet.add(sampExp);
							tissueExp.put(sample.getTissue(), hSet);
						}
						else {
							tissueExp.get(sample.getTissue()).add(sampExp); //otherwise add the sampExp to the existing HashSet
						}
					}
					ModelInputObj output = new ModelInputObj(); //finally set ModelInputObj according to above data
					output.setTissue_indivudual_expression(tissueExp);
					output.setVariant(variant);
					output.setTranscript(transcript);
					results.add(output);
				}
			}
			mmseqLine=mmseq.readLine();
		}
		return results;
	}
	

	
	public static void main (String[] args)
	{
		//CHANGE THESE
		String gtfPath = "C:\\Users\\Lucas\\Desktop\\demo\\demo\\gencode.v19.chr_patch_hapl_scaff.annotation.gtf.pseudo.retro.gtf";
		String mmseqPath = "C:\\Users\\Lucas\\Desktop\\demo\\demo\\brain.heart.adipose.lung.mmseq";
		String doePath = "C:\\Users\\Lucas\\Desktop\\demo\\demo\\gtex_doe_eur.regal.csv";
		String vcfPath = "C:\\Users\\Lucas\\Desktop\\demo\\demo\\GTEx_5M_plus_exome_185_Dec2012.phasing.noimputation.vcf.gz";
		String vcfIndexPath = "C:\\Users\\Lucas\\Desktop\\demo\\demo\\GTEx_5M_plus_exome_185_Dec2012.phasing.noimputation.vcf.gz.tbi";
		String covPath = "C:\\Users\\Lucas\\Desktop\\demo\\demo\\GTEx_5M_185_Dec2012.Eigenvectors.txt";
		String outputPath = "C:\\Users\\Lucas\\Desktop\\data.txt";
		
		InputObjGen i = new InputObjGen(gtfPath, mmseqPath, doePath, vcfPath, vcfIndexPath, covPath);
		HashSet<ModelInputObj> a = i.gen();
			File file = new File(outputPath);
			FileWriter fw = null;
			try {
				fw = new FileWriter(file.getAbsoluteFile());
			} catch (IOException e) {
				e.printStackTrace();
			}
			BufferedWriter bw = new BufferedWriter(fw); //create a writer to create an output txt file
		for(ModelInputObj m: a) //loop through all data
		{
			HashMap<Enum<Tissue>, HashSet<SampleExpression>> data = m.getTissue_indivudual_expression();
			for(HashSet<SampleExpression> hs: data.values())
			{
				for(SampleExpression samp : hs) //for every samp Exp
				{
					try {
						bw.write(samp.getIndividual().getVariant().getPos() //write data
						+ "\t" + m.getTranscript().getTranscriptID() + "\t"
						+ samp.getSample().getIndividual().getSubject_id() + "\t"
						 + samp.getMdiff() + "\t"
						 + samp.getVdiff() + "\t"
						 + samp.getUnique_hit_hap1() + "\t"
						 + samp.getUnique_hit_hap2() + "\t"
						 + samp.getIndividual().getGenotypeCode() + "\t"
						 + samp.getIndividual().getGenotype() + "\t"
						 + samp.getSample().getTissue()
						 );
						bw.newLine();
					} catch (IOException e) {
						e.printStackTrace();
					}
				}
			}
		}
		try {
			bw.close();
		} catch (IOException e) {
			e.printStackTrace();
		}
		System.out.println("done");
	}
}
